#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA

""" solve poisson equation in 2D and direchlet bc. """

from numpy import floor, linspace, meshgrid, roll, tile
from numpy import c_, asarray, ones, empty

from scipy.sparse.construct import spdiags
from scipy.sparse.linalg import spsolve as spsolve

from itertools import izip, count


class Poisson2DFiniteDifference(object):

    """
      Solve Poisson equation in 2D using a five points stencil, FD scheme.
      We assume Dirichlet boundary conditions on the grid.
      Only Dirichlet because that's what we need for calibrating
      Volatility.

      Carl Charrella, Mark Craddock and Nadima El-Hassan. The Calibration of stock option
      pricing models using inverse problem methodology

      Solve the problem

            -div( grad ( u(x,y) ) ) = f
            Dirichlet Boundary conditions.
    """
    __slots__ = [
        'm_dx',
        'm_dy',
        'm_xmin',
        'm_xmax',
        'm_ymin',
        'm_ymax',
        'm_nx_pts',
        'm_ny_pts',
        'm_xgrid',
        'm_ygrid',
        'm_xgposi',
        'm_ygposi',
        'm_left',
        'm_right',
        'm_bott',
        'm_top',
        'm_eng_ready',
        'm_ltype',
        'm_load_len',
        'm_vright',
        'm_vtop',
        'm_vbott',
        'm_vleft',
        'm_load',
        ]

    def __init__(self):
        """ instantiate object of type Poisson2D. """

        (self.m_dx, self.m_dy) = (None, None)
        (self.m_xmin, self.m_xmax) = (None, None)
        (self.m_ymin, self.m_ymax) = (None, None)

        (self.m_xgrid, self.m_ygrid) = (None, None)
        (self.m_xgposi, self.m_ygposi) = (None, None)
        (self.m_nx_pts, self.m_ny_pts) = (None, None)

        (self.m_bott, self.m_top) = (None, None)
        (self.m_left, self.m_right) = (None, None)
        (self.m_eng_ready, self.m_ltype) = (False, 'continuos')
        (self.m_load_len, self.m_vright) = (None, None)
        (self.m_vtop, self.m_vbott) = (None, None)
        (self.m_vleft, self.m_load) = (None, None)

    def set_xlxr(self, left, right):
        """ set left and right bc """

        (self.m_left, self.m_right) = (left, right)

    def set_bott_top(self, bott, top):
        """ set bottom and top bc """

        (self.m_bott, self.m_top) = (bott, top)

    def set_dx_dy(self, delta_x, delta_y):
        """ set grid delta x delta y """

        (self.m_dx, self.m_dy) = (delta_x, delta_y)

    def set_continuos_load(self, continuos_load):
        """ make load a function. """

        self.m_ltype = 'continuos'
        self.m_load = continuos_load

    def set_discrete_load(self, discrete_load):
        """ make load a matrix. Two different interfaces are used for
            clarity.
        """

        self.m_ltype = 'discrete'
        self.m_load = discrete_load

    def get_load(self):
        """ get load """

        load = self.__get_load()
        return (self.m_xgrid[1, -1], self.m_ygrid[-1, 1], load)

    def set_simu_domain(self, leftx, botty, rightx, topy):
        """ set boundary conditions """

        (self.m_xmin, self.m_xmax) = (float(leftx), float(rightx))
        (self.m_ymin, self.m_ymax) = (float(botty), float(topy))

    def reset_engine(self):
        """ reset pde solver """

        self.m_eng_ready = False

    def get_mesh(self):
        """ 2D grid use to solve the pde. """

        if self.m_xgrid is None or self.m_ygrid is None:
            self.set_up_engine()
        return (self.m_xgrid, self.m_ygrid)

    def get_indx_grid(self):
        """ grid to values dicrete load """

        if self.m_xgposi is None or self.m_ygposi is None:
            self.set_up_engine()

        return (self.m_xgposi, self.m_ygposi)

    def set_up_engine(self):
        """ get ready for solving the pde. """

        if self.m_eng_ready:
            return

        assert self.m_xmax > self.m_xmin
        assert self.m_ymax > self.m_ymin
        assert 0 < self.m_dx and 0 < self.m_dy

        nxsegs = int(floor((self.m_xmax - self.m_xmin) / self.m_dx))
        nysegs = int(floor((self.m_ymax - self.m_ymin) / self.m_dy))
        (self.m_nx_pts, self.m_ny_pts) = (max(nxsegs + 1, 3),
                max(nysegs + 1, 3))

        # adjust dx, dy
        self.m_dx = (self.m_xmax - self.m_xmin) / nxsegs
        self.m_dy = (self.m_ymax - self.m_ymin) / nysegs

        self.m_load_len = self.m_nx_pts * self.m_ny_pts - 2 \
            * (self.m_nx_pts + self.m_ny_pts - 2)

        self.m_xgrid = asarray(linspace(self.m_xmin, self.m_xmax,
                                self.m_nx_pts, 'True'))
        self.m_ygrid = asarray(linspace(self.m_ymin, self.m_ymax,
                                self.m_ny_pts, 'True'))

        (self.m_xgposi, self.m_ygposi) = (range(self.m_nx_pts),
                range(self.m_ny_pts))

        # Boundary conditions

        self.m_vbott = self.m_bott(c_[self.m_xgrid,
                                   ones(self.m_xgrid.shape)
                                   * self.m_ymin])
        self.m_vtop = self.m_top(c_[self.m_xgrid,
                                 ones(self.m_xgrid.shape)
                                 * self.m_ymax])
        self.m_vleft = self.m_left(c_[ones(self.m_ygrid.shape)
                                   * self.m_xmin, self.m_ygrid])
        self.m_vright = self.m_right(c_[ones(self.m_ygrid.shape)
                * self.m_xmax, self.m_ygrid])

        self.m_eng_ready = True

    def _eval_continuos_load(self, values):
        """ compute load using analytical formula """

        return self.m_load(values)

    def _eval_discrete_load(self, coor):
        """ values load represented as a matrix """

        return self.m_load[coor[0], coor[1]]

    def __get_load(self):
        """
            We are using lexicographic order to assemble the matrix.
            starting from the bottom left all the way up right.
            Assumes that the Image is lower-left
        """

        load = empty((self.m_load_len, 1), 'float')

        if self.m_ltype == 'continuos':
            (xgrid, ygrid) = meshgrid(self.m_xgrid[1:-1], self.m_ygrid[1:-1])
            for (i, val) in izip(count(), c_[xgrid.flatten(), ygrid.flatten()]):
                load[i] = self._eval_continuos_load(val)
        else:
            (xgrid, ygrid) = meshgrid(self.m_xgposi[1:-1], self.m_ygposi[1:-1])
            for (i, val) in izip(count(), c_[xgrid.flatten(), ygrid.flatten()]):
                load[i] = self._eval_discrete_load(val)
        return load

    def __assemble_rhs(self):
        """ compute rhs vector. """

        (nxpts, nypts) = (self.m_nx_pts - 2, self.m_ny_pts - 2)

        # Get Load
        rhs = self.__get_load()

        # Add boundary conditions

        (dx2, dy2) = (self.m_dx ** 2, self.m_dy ** 2)
        rhs[0:nxpts] += self.m_vbott[1:-1] / dy2
        rhs[-nxpts:] += self.m_vtop[1:-1] / dy2

        indx = nxpts * asarray(range(0, nypts), 'int')
        rhs[indx] += self.m_vleft[1:-1] / dx2
        rhs[-(indx[::-1] + 1)] += self.m_vright[1:-1] / dx2

        return rhs

    def __assemble_matrix(self):
        """ assemble matrix """

        rhs_len = self.m_load_len
        (dx2, dy2) = (self.m_dx ** 2, self.m_dy ** 2)
        (nxpts, nypts) = (int(self.m_nx_pts - 2), int(self.m_ny_pts
                      - 2))

        # Main Diagonal

        shape = (rhs_len, 1)
        diag = 2. * ones(shape) * (1. / dx2 + 1. / dy2)

        # Off Diagonals

        pattx = ones(nxpts)
        pattx[-1] = 0
        diag_1 = -tile(pattx, nypts) / dx2  # lower
        diag_nx = -ones(shape) / dy2  # lower
        diag1 = roll(diag_1, 1)  # upper
        diagnx = -ones(shape) / dy2  # upper

        # Assemble

        diags = [0, -1, -nxpts, 1, nxpts]
        data = asarray([list(diag.flatten()), list(diag_1.flatten()),
                       list(diag_nx.flatten()), list(diag1.flatten()),
                       list(diagnx.flatten())])

        return spdiags(data, diags, rhs_len, rhs_len).tocsc()

    def discretize(self):
        """ discretize pde """

        self.set_up_engine()
        return (self.__assemble_matrix(), self.__assemble_rhs())

    def get_rhs(self):
        """ get right hand side """

        return self.__assemble_rhs()

    def get_discretized_bc(self):
        """ returns dirichlet bc as vectors. """

        return (self.m_vbott, self.m_vtop, self.m_vleft, self.m_vright)


    @classmethod
    def post_process(cls, nx_opts, ny_opts, vbott, vtop, vleft, vright, vsol):
        """ post process results """

        solu = empty((ny_opts, nx_opts))

        (solu[-1, :], solu[0, :]) = (vbott.flatten(), vtop.flatten())
        (solu[:, 0], solu[:, -1]) = (vleft.flatten(), vright.flatten())

        (nxpts, nypts) = (int(nx_opts - 2), int(ny_opts - 2))
        tmp = vsol.reshape((nypts, nxpts))
        for posi in range(0, nypts):
            solu[-(posi + 2), 1:-1] = tmp[posi, :]

        return solu


    def solve(self):
        """ discretize and solve pde """

        (stiff_mat, rhs) = self.discretize()

        tmp = spsolve(stiff_mat, rhs)
        solu = self.post_process(self.m_nx_pts, self.m_ny_pts,
                                self.m_vbott, self.m_vtop,
                                self.m_vleft[::-1], self.m_vright[::-1],
                                tmp)

        return (self.m_xgrid, self.m_ygrid, solu)

